# -*- coding: utf-8 -*-
"""
Created on Sat Feb 20 18:42:13 2021
this code is modified from mnist_cnn for detecting cat and dog
@author: LI
"""


import tensorflow as tf
import os
import sys
import matplotlib.pyplot as plt
import numpy as np
import glob
from skimage import io, transform
import cv2 as cv
import copy
from tensorflow.python.platform import gfile


# 载入数据集
path = '/home/ljl/Desktop/save_roi/dst5-x6-269/'
row = 28
col = 28
c = 1
img_type = '.png'
class_num = 10
# 读取数据


def kind_label(idx):
    dict_label = {
        0:"up",
        1:"pingtai",
        2:"tiegui",
        3:"dimian",
        4:"filters",
        5:"remain"
        }
    return dict_label[idx]



def test_readImgs(path):
    cate = [path + x for x in os.listdir(path) if os.path.isdir(path + x)]
    img_org = []
    img_ = []
    label = []
    np_label = np.zeros((1, class_num), dtype=np.int)
    org_label = list(np_label[0])
    for idx, folder in enumerate(cate):  # 0  4   2   3   1
        for im in glob.glob(folder + '/*' + img_type):
            if not os.path.exists(im):
                print(">>>>>>>>>  file is not exit   <<<<<<<<<<<<\n")
                sys.exit(0)
            img = cv.imread(im)
            img_org.append(img)

            test_img = transform.resize(img, (row, col, c), mode='constant')

            img_.append(test_img)

            make_label = copy.deepcopy(org_label)
            #make_label[idx] = 1
            label.append(make_label)

    return np.asarray(img_, np.float32), np.asarray(label, np.int32), img_org

#load  model
model_path = './saved-model'
model_path_pb = './models'
# 加载模型
with tf.Session(graph=tf.Graph())as sess:

    # tf.saved_model.loader.load(sess,[tf.saved_model.tag_constants.TRAINING],save_path)
    sess.run(tf.global_variables_initializer())  # 加了这句话就会导致图初始化

    # 加载模型
    if 0:       #cpkt
        saver = tf.train.import_meta_graph(model_path + '/model.ckpt.meta')  # 先加载meta文件，具体到文件名
        saver.restore(sess, tf.train.latest_checkpoint(model_path))  # 加载检查点文件checkpoint，具体到文件夹即可
        graph = tf.get_default_graph()  # 绘制tensorflow图
        input_x = sess.graph.get_tensor_by_name('x-input:0')
        input_y = sess.graph.get_tensor_by_name('y-input:0')
        keep_prob = sess.graph.get_tensor_by_name('keep_prob:0')
        output = sess.graph.get_tensor_by_name('softmax:0') 
    else:       #pb
        with gfile.FastGFile(model_path_pb + '/model.pb', 'rb') as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
            sess.graph.as_default()
            tf.import_graph_def(graph_def, name='') # 导入计算图
              

            input_x = sess.graph.get_tensor_by_name('x-input:0')
            keep_prob = sess.graph.get_tensor_by_name('keep_prob:0')
            output = sess.graph.get_tensor_by_name('softmax:0')
    
    path = '/home/ljl/Desktop/save_roi/dst5-x6-269/'
    #save_process_roi = "/home/ljl/Desktop/save_roi/dst5-x6-472 -result/"
    save_process_roi = "/home/ljl/Desktop/save_roi/dst5-x6-269-result/"
    #save_process_roi = "/home/ljl/Desktop/save_roi/dst5-x6-674 -result/"

    test_img, test_label, test_show = test_readImgs(path)
    for i in range(len(test_label)):
        prediction = sess.run(output, feed_dict={input_x: test_img[i:i+1], keep_prob: 1.0})

        pre_idx = np.argmax(prediction[0])
        print(prediction[0].shape)
        print('Training prediction=', prediction)
        print(' -> output class is: ', pre_idx + 1, "     is:  ",kind_label(1),  "\n",
              '-> prediction   is: ', prediction[0][pre_idx])
        print('>>>>>>>>>>>>>>>\t end \t<<<<<<<<<<<<<')
        
        if (pre_idx+1 < 5 and  prediction[0][pre_idx] > 0.85) or (pre_idx + 1 == 5 and prediction[0][pre_idx] > 0.95):
            print("*************delete")
            cv.imshow('aa', test_show[i])
            cv.waitKey(0)
            cv.destroyAllWindows()
            continue
        else:       #4
            cv.imwrite(save_process_roi + str(i) + ".png",test_show[i])
            
            cv.imshow('aa', test_show[i])
            cv.waitKey(0)
            cv.destroyAllWindows()

    # 0  4   2   3   1
